import os
osPath = os.path
import sys
import time
import logging
import threading
import collections

from random import randint
from cPickle import load, dump
from functools import wraps

import Image

def locked(lock):
	"""
	Synchronization decorator.
	"""
	def decorated(f):
		@wraps(f)
		def lockedFunction(*args, **kw):
			with lock:
				return f(*args, **kw)
		return lockedFunction
	return decorated

class ImageHistory(object):
	"""
	A brief history of previously used images
	"""
	CacheLock = threading.RLock()
	def __init__(self):
		self.needsWrite = False
		self.priorWalls = set()
		self.priorWalls = set(self.loadWalls())
		self.count = 0

	def loadWalls(self):
		if not self.priorWalls and osPath.exists('priorWalls.txt'):
			with open('priorWalls.txt', 'rb') as fp:
				return set(f.strip() for f in fp.readlines())
		return set()

	@locked(CacheLock)
	def addWall(self, wall):
		self.needsWrite = True
		self.priorWalls.add(wall)

	@locked(CacheLock)
	def writeWalls(self):
		if self.needsWrite and not self.count:
			with open('priorWalls.txt', 'wb') as fp:
				existing = (w for w in self.priorWalls if osPath.exists(w))
				fp.write('\n'.join(existing))
				fp.flush()
			self.needsWrite = False

	@locked(CacheLock)
	def removeWalls(self, walls):
		if walls:
			self.priorWalls -= set(walls)
			self.needsWrite = True

	def getAvailable(self, files, dontWant = set()):
		return list(set(files) - self.priorWalls - dontWant)


	@locked(CacheLock)
	def bump(self):
		self.count += 1

	@locked(CacheLock)
	def unBump(self):
		self.count -= 1
		
	def __enter__(self):
		self.bump()
		return self

	def __exit__(self, exc_type, exc_value, traceback):
		self.unBump()

# For quick filtering
ImageExtensions = set((
	'.png', '.jpg', '.jpeg', '.gif', '.fli', '.flc', '.fpx', '.gbr', '.gd',
	'.ico', '.im', '.pcd', '.pcx', '.ppm', '.psd', '.sgi', '.tga', '.tiff',
	'.wal', '.xbm', '.xpm', '.bmp',
	))

class DirEntry(object):
	__slots__ = ('path', 'stat', '_files', 'files')
	def __init__(self, path, files = None, stat = os.stat):
		self.path = path
		self.stat = stat(path)
		self._files = files

	@property
	def files(self):
		"""
		Delay loading of files until we care
		"""
		if self._files is None:
			self._files = os.walk(self.path).next()
		return self._files

class OnFileCache(object):
	def __init__(self):
		self.DirCache = {}
		self.history = ImageHistory()
		self.changed = False
		if osPath.exists('dircache.txt'):
			try:
				self.DirCache = load(open('dircache.txt', 'rb'))
			except:
				pass

	def get(self, dire):
		return self.DirCache.get(dire)

	def set(self, dire, files):
		self.changed = True
		self.DirCache[dire] = DirEntry(dire, files)
		return dire

	def update(self, dire, files = None, stat = os.stat):
		if dire in FileCache.DirCache and self.DirCache[dire].stat == stat(dire):
			return dire

		return self.set(dire, files)

	def flush(self):
		if self.changed:
			dump(self.DirCache, open('dircache.txt', 'wb'), -1)
			self.changed = False

FileCache = OnFileCache()
update = FileCache.update

def flushWalls():
	FileCache.history.writeWalls()
	
def expandDir_t(dire, nDe):
	"""
	Expand dirs in a thread.

	:param dire: Root directory
	:param nDe: Callback
	"""
	dirM = collections.deque()
	append = dirM.append

	def _updater(append, root):
		for root, dirs, files in os.walk(root):
			if files:
				update(root, files)
				append(root)

	try:
		root, dirs, files = os.walk(dire).next()
	
		threads = [threading.Thread(target = _updater,
									args = (append, osPath.join(root, d)))
				   for d in dirs]

		for t in threads:
			t.start()

		if files:
			update(root, files)
			append(root)

		for t in threads:
			t.join()
	except StopIteration:
		pass
		
	nDe((dire,dirM))

def expandDirs(dirs):
	"""
	Find all dirs, recurse down any path starting with '+'

	:param dirs: A list of dirs
	"""
	newDirs = (FileCache.update(d) for d in dirs if d[0] != '+' and osPath.exists(d))
	newDirs = [(d,(d,)) for d in newDirs if d]

	nDe = newDirs.append

	dThreads = [threading.Thread(target = expandDir_t,
								args = (dire, nDe))
			   for dire in (d[1:] for d in dirs if d[0] == '+' and osPath.exists(d[1:]))]

	# Start Walking
	for t in dThreads:
		t.start()

	# Wait for all walkers to finish
	for t in dThreads:
		t.join()
		
	FileCache.flush()
	return newDirs

def chooseDir(dirs):
	"""
	Choose a directory from a tree
	"""
	nDirs = len(dirs)
	dirs = dirs[randint(0, nDirs - 1)][1]
	nDirs = len(dirs)
	if nDirs:
		path = dirs[randint(0, nDirs - 1)]
		return path

def getNewImage(directory, dontWant = set()):
	"""
	Get a new image from a directory

	:param directory: A directory full of files
	"""
	opj = osPath.join
	ops = osPath.splitext

	files = FileCache.get(directory).files
	
	if files:
		files = [opj(directory, f) for f in files if ops(f)[1] in ImageExtensions]
		with FileCache.history as cache:
			availableFiles = cache.getAvailable(files, dontWant)
			
			if not availableFiles:
				cache.removeWalls(files)
				availableFiles = cache.getAvailable(files, dontWant)
			
			tries = 0
			while tries < 3 and availableFiles:
				idx = randint(0, len(availableFiles) - 1)
				filename = availableFiles[idx]
				cache.addWall(filename)
				try:
					image = Image.open(filename)
					image.draft('RGBA', image.size)
					image.info['filename'] = filename
					return image
				except Exception, e:
					availableFiles.remove(filename)
				tries += 1
	return None
